package com.tjbklx1.zkclient.queue;

import java.util.Collections;
import java.util.Comparator;
import java.util.List;

import org.I0Itec.zkclient.ExceptionUtil;
import org.I0Itec.zkclient.ZkClient;
import org.I0Itec.zkclient.exception.ZkNoNodeException;

/**
 * 简单的分布式队列
 * @author Administrator
 *
 * @param <T>
 */
public class DistributedSimpleQueue<T> {

	protected final ZkClient zkClient;
	protected final String root;

	// 顺序节点的名称
	protected static final String Node_NAME = "n_";

	public DistributedSimpleQueue(ZkClient zkClient, String root) {
		this.zkClient = zkClient;
		this.root = root;
	}

	public int size() {
		return zkClient.getChildren(root).size();
	}

	public boolean isEmpty() {
		return zkClient.getChildren(root).size() == 0;
	}

	/**
	 * 向队列提交数据
	 * 
	 * @param element
	 * @return
	 * @throws Exception
	 */
	public boolean offer(T element) throws Exception {
		//完整路径
		String nodeFullPath = root.concat("/").concat(Node_NAME);
		try {
			//创建顺序节点
			zkClient.createPersistentSequential(nodeFullPath, element);
		} catch (ZkNoNodeException e) {
			//如果创建顺序节点失败，就要先创建根节点，再创建顺序节点
			zkClient.createPersistent(root);
			offer(element);
		} catch (Exception e) {
			throw ExceptionUtil.convertToRuntimeException(e);
		}
		return true;
	}

	@SuppressWarnings("unchecked")
	public T poll() throws Exception {
		try {
			//查询所有的子节点
			List<String> list = zkClient.getChildren(root);
			if (list.size() == 0) {
				return null;
			}
			//排序
			Collections.sort(list, new Comparator<String>() {
				public int compare(String lhs, String rhs) {
					return getNodeNumber(lhs, Node_NAME).compareTo(
							getNodeNumber(rhs, Node_NAME));
				}
			});

			//按顺序读取子节点的数据,并删除,返回节点数据 (poll方法的核心: zkClient.delete 后返回 )
			for (String nodeName : list) {
				String nodeFullPath = root.concat("/").concat(nodeName);
				try {
					T node = (T) zkClient.readData(nodeFullPath);
					zkClient.delete(nodeFullPath);
					return node;
				} catch (ZkNoNodeException e) {
					// ignore
				}
			}
			return null;
		} catch (Exception e) {
			throw ExceptionUtil.convertToRuntimeException(e);
		}
	}

	private String getNodeNumber(String str, String nodeName) {
		int index = str.lastIndexOf(nodeName);
		if (index >= 0) {
			index += Node_NAME.length();
			return index <= str.length() ? str.substring(index) : "";
		}
		return str;
	}

}
